##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(ArchR)
library(Seurat)
library(BSgenome.Hsapiens.UCSC.hg38)
library(GenomicFeatures)
library(tibble)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--rna_data_file"), type = "character"),
    make_option(c("--motif_annotation_file"), type = "character"),
    make_option(c("--gff3_file"), type = "character"),
    make_option(c("--gene_ensg_file"), type = "character"),
    make_option(c("--peak_type"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){

    ## 单细胞的atac
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 表达文件
    rna_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined.annotationCellType.qc.Rdata"

    ## motif所在peak的注释
    motif_annotation_file <- "~/20231121_singleMuti/results/celltype_plot/diff_peak/germ/AllPeak.containMotif.tsv"

    ## gtf_file 
    gff3_file <- "~/ref/GTF/gencode.v32.annotation.gff3"
    gene_ensg_file <- "~/ref/GTF/gencode.v32.gene_ensg.tsv"

    ## 选取调控的peak类型
    peak_type <- "positive"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/peak2gene/germ_mfuzz_new"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
rna_data_file <- opt$rna_data_file
motif_annotation_file <- opt$motif_annotation_file
gff3_file <- opt$gff3_file
gene_ensg_file <- opt$gene_ensg_file
peak_type <- opt$peak_type
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
a <- load(rna_data_file)
b <- load(comine_data_file)
motif_annotation <- fread(motif_annotation_file)

corrCutoff <- 0.5       # Default in plotPeak2GeneHeatmap is 0.45

##########################################################################################
## 以防表达矩阵被修改
## atac加入表达矩阵
txdb <- makeTxDbFromGFF(gff3_file)
gene_ensg <- read.table(gene_ensg_file)

gene_gr <- genes(txdb, columns = c("GENEID"))
colnames(gene_ensg) <- c("GENEID" , "name")

gr_df <- data.frame(
  seqnames = seqnames(gene_gr),
  start = start(gene_gr),
  end = end(gene_gr),
  strand = strand(gene_gr),
  GENEID = names(mcols(gene_gr)$GENEID)
)

## gene name对应多个ensg的，选第一个
gr_df <- merge( gr_df , gene_ensg , by = "GENEID" )
gr_df <- gr_df %>% 
dplyr::group_by( name ) %>%
dplyr::summarize( seqnames = seqnames[1] , start = start[1] , end = end[1] , strand = "*"  )

gr <- GRanges(
  seqnames = Rle(gr_df$seqnames),
  ranges = IRanges(start = gr_df$start, end = gr_df$end),
  strand = gr_df$strand,
  name = gr_df$name
)

names(gr) <- gr_df$name

use_gene <- rownames(scrnat@assays$RNA$counts)[rownames(scrnat@assays$RNA$counts) %in% names(gr)]
gr <- gr[use_gene]

seruat_cds_all <- SummarizedExperiment(scrnat@assays$RNA$counts[use_gene,] , rowRanges = gr)
colnames(seruat_cds_all) <- gsub( "_" , "#" ,  colnames(seruat_cds_all) )

## 整合atac和表达
testis_combined_peak_combineRNA <- addGeneExpressionMatrix(input = testis_combined_peak_combineRNA, seRNA = seruat_cds_all, force = TRUE)
testis_combined_peak_combineRNA <- addImputeWeights(testis_combined_peak_combineRNA)

rm(scrnat)

##########################################################################################
## To identify peak-to-gene links in ArchR
projHeme5 <- testis_combined_peak_combineRNA

## 保证peak用之前的，后面用peak的方式会修改原始peak对应的参数
projHeme5 <- addPeakSet(projHeme5,projHeme5@peakSet,force = TRUE)
projHeme5 <- addPeakMatrix(projHeme5)
projHeme5 <- addPeak2GeneLinks(ArchRProj = projHeme5 , useMatrix = "GeneExpressionMatrix")

p2gMat <- plotPeak2GeneHeatmap(
  projHeme5, 
  corCutOff = corrCutoff, 
  groupBy="cell_type",
  nPlot = 1000000, returnMatrices=TRUE, 
  k=25, seed=1)

tmp_motif_annotation <- motif_annotation %>% 
  dplyr::group_by(peakName) %>% 
  dplyr::summarize( motif=paste0(motif_in , collapse = "|") , motif_num = length(unique(motif_in)) )

comb2 <- merge( p2gMat$Peak2GeneLinks , tmp_motif_annotation , by.x = "peak" , by.y = "peakName" )

out_file <- paste0(out_path , "/peak-gene_all_postive.tsv")
write.table( comb2 , out_file , row.names = F , sep = "\t" , quote = F )

rm(motif_annotation)

if(1!=1){
  ## 提取显著正相关和负相关的peak
  p2g <- getPeak2GeneLinks(
    ArchRProj = projHeme5,
    corCutOff = -1,
    resolution = 1000000000,
    FDRCutOff = 1e-04,
    varCutOffATAC = 0.25,
    varCutOffRNA = 0.25, 
    returnLoops = FALSE
  )

  mATAC <- readRDS(metadata(p2g)$seATAC)[p2g$idxATAC, ]
  mRNA <- readRDS(metadata(p2g)$seRNA)[p2g$idxRNA, ]
  p2g$peak <- paste0(rowRanges(mATAC))
  p2g$gene <- rowData(mRNA)$name
  comb <- p2g
  comb$type <- ifelse( comb$Correlation > corrCutoff, "positive" ,  "other" )
  comb <- subset( comb , type != "" )

  ## peak所在motif的注释
  tmp_motif_annotation <- motif_annotation %>% 
  dplyr::group_by(peakName) %>% 
  dplyr::summarize( motif=paste0(motif_in , collapse = "|") , motif_num = length(unique(motif_in)) )

  comb2 <- merge( comb , tmp_motif_annotation , by.x = "peak" , by.y = "peakName" )

  ## 所有显著相关的peak-gene，不考虑相关系数
  out_file <- paste0(out_path , "/peak-gene_all_postive-negative.tsv")
  write.table( comb2 , out_file , row.names = F , sep = "\t" , quote = F )

  rm(motif_annotation)
}


###########################################################################################
## 选取感兴趣的peak集合

tmp_peak <- data.frame(projHeme5@peakSet)
tmp_peak <- paste0(tmp_peak$seqnames , ":" , tmp_peak$start , "-" , tmp_peak$end)

if( peak_type == "positive" ){
  use_peak <- unique(comb2$peak)
}else if( peak_type == "other" ){
  use_peak <- tmp_peak[!tmp_peak %in% unique(comb2$peak)]
}

tmp_peak_index <- which(tmp_peak %in% use_peak)

###########################################################################################
## peak只用存在peak-gene关系的peak
atac_proj <- projHeme5
atac_proj@peakSet <- atac_proj@peakSet[tmp_peak_index,]
## 使用自己定义的peak集合
atac_proj <- addPeakSet(atac_proj,atac_proj@peakSet,force = TRUE)
atac_proj <- addPeakMatrix(atac_proj)

## 重新计算这部分peak上的motif活性，背景用peak-gene的peak
atac_proj <- addMotifAnnotations(ArchRProj = atac_proj, motifSet = "cisbp", name = "Motif",force = TRUE)
atac_proj <- addBgdPeaks(atac_proj , force = TRUE)
atac_proj <- addDeviationsMatrix(
  ArchRProj = atac_proj, 
  peakAnnotation = "Motif",
  force = TRUE
)

corGIM_MM <- correlateMatrices(
    ArchRProj = atac_proj,
    useMatrix1 = "GeneExpressionMatrix",
    useMatrix2 = "MotifMatrix" 
)

corGIM_MM <- corGIM_MM[corGIM_MM$MotifMatrix_matchName == corGIM_MM$GeneExpressionMatrix_name,]
corGIM_MM <- corGIM_MM[!is.na(corGIM_MM$cor),]
## 输出
out_file <- paste0( out_path , "/cor.motif_atac-rna." , peak_type , ".tsv" )
write.table( corGIM_MM , out_file , row.names = F , sep = "\t" )

###########################################################################################
## motif
motif_se <- getMatrixFromProject(atac_proj, useMatrix="MotifMatrix")
motif_mat <- assays(motif_se)$deviation
## 转化为数值矩阵
motif_mat_num <- apply(motif_mat , 1 , as.numeric)
rownames(motif_mat_num) <- colnames(motif_mat)
motif_mat_num <- t(motif_mat_num)
## 去除ENSG基因
#motif_mat_num <- motif_mat_num[grep( "ENSG" , rownames(motif_mat_num) , invert = T),]
## 去除没有表达的motif
motif_mat_num <- motif_mat_num[corGIM_MM$MotifMatrix_name,]
## 去除名字的_
rownames(motif_mat_num) <- sapply(strsplit(rownames(motif_mat_num) , "_") , "[" , 1)

cell_cluster <- data.frame( cell = rownames(atac_proj@cellColData) , cluster = atac_proj@cellColData$seurat_clusters )
cell_cluster$use <- paste0( cell_cluster$cell , "_" , cell_cluster$cluster )
rownames(cell_cluster) <- cell_cluster$cell

## 用于计算motif的评价表达
motif_mat_num <- motif_mat_num[,rownames(atac_proj@cellColData)]
GSM_mat_num <- motif_mat_num

#colnames(motif_mat_num) <- cell_cluster[colnames(motif_mat_num),"use"]
out_file <- paste0(out_path , "/Motif.Peak-Gene." , peak_type , ".rds")
saveRDS( GSM_mat_num , out_file )

## 计算motif在每个细胞的评价表达
sco_motif <- sapply(unique(unique(atac_proj@cellColData$cell_type)),function(x){
    print(x)
    sapply(unique(rownames(GSM_mat_num)),function(y){
        mean(
            as.numeric(
                as.vector(GSM_mat_num[y,which(atac_proj@cellColData$cell_type==x)])
            )
        )
    })
})

sco_motif <- data.frame(sco_motif)
sco_motif$gene <- rownames(sco_motif)
sco_motif <- sco_motif[,c( ncol(sco_motif) , (1:ncol(sco_motif)-1) )]

out_file <- paste0(out_path , "/Motif.Peak-Gene.MeanByCellType." , peak_type , ".tsv")
write.table( sco_motif , out_file , row.names = F , quote = F , sep = "\t" )

###########################################################################################
## 保证peak用之前的，后面用peak的方式会修改原始peak对应的存贮
projHeme5 <- addPeakSet(projHeme5,projHeme5@peakSet,force = TRUE)
projHeme5 <- addPeakMatrix(projHeme5)
projHeme5 <- addMotifAnnotations(ArchRProj = projHeme5, motifSet = "cisbp", name = "Motif",force = TRUE)
projHeme5 <- addBgdPeaks(projHeme5 , force = TRUE)
projHeme5 <- addDeviationsMatrix(
  ArchRProj = projHeme5, 
  peakAnnotation = "Motif",
  force = TRUE
)

projHeme5 <- addPeak2GeneLinks(ArchRProj = projHeme5 , useMatrix = "GeneExpressionMatrix")
